import asyncio
import cv2
import numpy as np
import websockets
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, RTCIceCandidate
from av import VideoFrame
import fractions
import json

class CustomVideoStreamTrack(VideoStreamTrack):
    def __init__(self, rtsp_url):
        super().__init__()
        self.cap = cv2.VideoCapture(rtsp_url)

        self.frame_count = 0

    async def recv(self):
        self.frame_count += 1
        ret, frame = self.cap.read()
        if not ret or not self.cap.isOpened():
            print("无法从RTSP流中读取帧")
            return None
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        video_frame = VideoFrame.from_ndarray(frame, format="rgb24")
        video_frame.pts = self.frame_count
        video_frame.time_base = fractions.Fraction(1, 30)
        return video_frame

async def on_ice_candidate(pc, candidate):
    print(f"ICE 候选人: {candidate}")

async def handle_peer(pc, ws, offer):
    @pc.on("icecandidate")
    async def on_icecandidate(candidate):
        if candidate:
            await ws.send(json.dumps({
                'candidate': candidate.to_sdp(),
                'peer_id': pc.id
            }))

    await pc.setLocalDescription(offer)
    await ws.send(json.dumps({
        'sdp': pc.localDescription.sdp,
        'type': pc.localDescription.type,
        'peer_id': pc.id
    }))

    answer = await ws.recv()
    answer_dict = json.loads(answer)
    remote_desc = RTCSessionDescription(sdp=answer_dict['sdp'], type=answer_dict['type'])
    await pc.setRemoteDescription(remote_desc)

async def setup_webrtc_and_run(ws_url, rtsp_url):
    video_sender = CustomVideoStreamTrack(rtsp_url)
    pcs = {}  # 使用字典来跟踪所有PeerConnections

    async with websockets.connect(ws_url) as ws:
        # 首先发送自身的peer_id给信令服务器
        peer_id = "producer_id"  # 这里可以是一个唯一标识符，例如UUID
        await ws.send(json.dumps({'peer_id': peer_id}))

        while True:
            msg = await ws.recv()
            msg_dict = json.loads(msg)
            print("读取到消息：",msg_dict)
            # if msg_dict['type'] == 'offer':
            if msg_dict['type'] == 'p2p':
                # offer = RTCSessionDescription(sdp=msg_dict['sdp'], type=msg_dict['type'])

                # 为每个新的对等方创建一个新的PeerConnection
                pc = RTCPeerConnection()

                pc.id = peer_id  # 假设每个消息都有一个唯一的peer_id
                pc.addTrack(video_sender)
                pcs[peer_id] = pc
                offer = await pc.createOffer()
                await handle_peer(pc, ws, offer)

            elif msg_dict['type'] == 'icecandidate':
                peer_id = msg_dict['peer_id']
                candidate = RTCIceCandidate.from_sdp(msg_dict['candidate'])
                await pcs[peer_id].addIceCandidate(candidate)

async def main():
    ws_url = "ws://localhost:8765"  # WebSocket URL of the signaling server
    rtsp_url = 'rtsp://admin:123456@192.168.0.110:554/h264/ch1/main/av_stream'
    await setup_webrtc_and_run(ws_url, rtsp_url)

if __name__ == "__main__":
    asyncio.run(main())
